from transformers import ViTFeatureExtractor, ViTModel
from PIL import Image
import numpy as np 
import torch

class VitBase():

    def __init__(self):
        model_path = r"E:\nlp\pretrain_models\image\vit-base-patch16-224-in21k"
        self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_path)
        self.model = ViTModel.from_pretrained(model_path)
    
    def extract_feature(self, imgs):
        inputs = self.feature_extractor(images=imgs, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
        last_hidden_states =  outputs.last_hidden_state
        features = np.reshape(last_hidden_states.numpy(),(len(imgs),-1))
        return features
